/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 * Description:
 * Author: l30050806
 * Create: 2024-12-27
 */

#include "transport_ipc_mem.h"
#include "log.h"
#include "network_manager_pub.h"

namespace hccl {
using namespace std;
TransportIpcMem::TransportIpcMem(const std::unique_ptr<NotifyPool> &notifyPool, const HcclNetDevCtx &netDevCtx,
    const HcclDispatcher &dispatcher, u32 localRankId, u32 remoteRankId)
    : TransportMem(notifyPool, netDevCtx, dispatcher, localRankId, remoteRankId)
{}

TransportIpcMem::~TransportIpcMem()
{}

HcclResult TransportIpcMem::SetSocket(const std::shared_ptr<HcclSocket> &socket)
{
    HCCL_INFO("TransportIpcMem doesn't need to set socket");
    return HCCL_SUCCESS;
}

HcclResult TransportIpcMem::Connect(s32 timeoutSec)
{
    HCCL_INFO("TransportIpcMem doesn't need to connect socket");
    return HCCL_SUCCESS;
}

HcclResult TransportIpcMem::Write(const RmaBufferSlice &remoteRmaBufferSlice, const RmaBufferSlice &localRmaBufferSlice, const rtStream_t &stream)
{
    CHK_RET(TransportIpc(remoteRmaBufferSlice, localRmaBufferSlice, stream));
    return HCCL_SUCCESS;
}

HcclResult TransportIpcMem::Read(const RmaBufferSlice &localRmaBufferSlice, const RmaBufferSlice &remoteRmaBufferSlice, const rtStream_t &stream)
{
    CHK_RET(TransportIpc(localRmaBufferSlice, remoteRmaBufferSlice, stream));
    return HCCL_SUCCESS;
}

HcclResult TransportIpcMem::AddOpFence(const rtStream_t &stream)
{
    HCCL_INFO("TransportIpcMem doesn't need to add op fence");
    return HCCL_SUCCESS;
}

HcclResult TransportIpcMem::TransportIpc(
    const RmaBufferSlice &dstRmaBufferSlice, const RmaBufferSlice &srcRmaBufferSlice, const rtStream_t &stream)
{
    CHK_PTR_NULL(dstRmaBufferSlice.addr);
    CHK_PTR_NULL(srcRmaBufferSlice.addr);
    Stream hcclStream(stream);
    if (dstRmaBufferSlice.memType == HcclMemType::HCCL_MEM_TYPE_HOST ||
        srcRmaBufferSlice.memType == HcclMemType::HCCL_MEM_TYPE_HOST) {
        CHK_RET(dispatcher_->MemcpyAsyncWithoutCheckKind(dstRmaBufferSlice.addr, dstRmaBufferSlice.len, srcRmaBufferSlice.addr,
            srcRmaBufferSlice.len, HcclRtMemcpyKind::HCCL_RT_MEMCPY_KIND_DEVICE_TO_DEVICE, hcclStream, remoteRankId_,
            hccl::LinkType::LINK_HCCS));
    } else {
        CHK_RET(dispatcher_->MemcpyAsync(dstRmaBufferSlice.addr, dstRmaBufferSlice.len, srcRmaBufferSlice.addr,
            srcRmaBufferSlice.len, HcclRtMemcpyKind::HCCL_RT_MEMCPY_KIND_DEVICE_TO_DEVICE, hcclStream, remoteRankId_,
            hccl::LinkType::LINK_HCCS));
    }
    return HCCL_SUCCESS;
}
}  // namespace hccl
